import torch

x_data = torch.Tensor([[1.0], [2.0], [3.0]])#3*1的矩阵
y_data = torch.Tensor([[2.0], [4.0], [6.0]])#3*1的矩阵


class LinerModelWithTorch(torch.nn.Module):  # 继承torch.nn.Module
    def __init__(self):
        super(LinerModelWithTorch, self).__init__()
        self.linear = torch.nn.Linear(1, 1)  # 两个参数分别是输入样本特征的维度和输出样本特征的维度

    # 已经继承父类，即继承了Module的callable函数，callable函数中调用了forward，我们进行重写
    def forward(self, x):  # 要重写父类的forward函数
        y_pred = self.linear(x)  # 用的是liner的callable函数
        return y_pred


model = LinerModelWithTorch()

criterion = torch.nn.MSELoss(size_average=False)  # 构造损失函数，能得到yhat和y  size_average=False不使用均方误差

# 我们的linerModel继承了父类，父类中有一个成员函数，它会检查module里的所有成员，如果该成员有相应的权重，它会把该成员加到model.parameters()的参数集合上
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 第一个参数是告诉优化器需要优化哪些参数，第二个是学习率

# 进行训练
for epoch in range(1000):#100次epoch
    y_pred = model(x_data)#获取yhat
    loss = criterion(y_pred,y_data)#依据20行的设置计算损失
    print(epoch,loss.item())

    optimizer.zero_grad()#梯度清零
    loss.backward()#反向传播
    optimizer.step()#优化器更新，依据23行的设置更新所有参数

#Output weight and bias
print('w= ',model.linear.weight.item())
print('b= ',model.linear.bias.item())

#TestModel
x_test = torch.Tensor([[4.0]])
y_test = model(x_test)
print('y_pred = ',y_test.data)